%{
#include <iostream>
#include <string>
#include <vector>
#include <memory>
#include "ASTNode.h"
#include "SymbolTable.h"
#include "CodeGenerator.h"
#include "Optimizer.h"

using namespace std;

extern int yylex();
extern int yyparse();
extern FILE *yyin;

void yyerror(const char *s) {
    extern char *yytext;
    extern int yylineno;
    
    cerr << "Error at line " << yylineno << ": " << s;
    if (yytext) cerr << " near '" << yytext << "'";
    cerr << endl;
}

ASTNode* root = nullptr;

%}

%code requires {
#include "ASTNode.h"
}

%union {
    ASTNode* node;       // 所有 AST 节点
    std::string* str;    // 标识符和字符串字面量
    int num;             // 数值字面量
}

%token <num> NUMBER
%token <str> ID
%token INT VOID
%token IF ELSE WHILE BREAK CONTINUE RETURN
%token ADD SUB MUL DIV MOD NOT AND OR
%token LT GT LE GE EQ NE ASSIGN
%token SEMICOLON COMMA LPAREN RPAREN LBRACE RBRACE

%type <node> CompUnit FuncDef Param Stmt Block Expr
%type <node> LOrExpr LAndExpr RelExpr AddExpr MulExpr UnaryExpr PrimaryExpr
%type <node> type param_list Stmt_list arg_list

%start CompUnit

%%

CompUnit: FuncDef 
        { 
            $$ = new ASTNode("CompUnit"); 
            $$->addChild($1); 
            root = $$;  // 确保 root 指向 CompUnit
        }
        | CompUnit FuncDef
        { 
            $1->addChild($2); 
            $$ = $1; 
            root = $$;  // 更新 root 保持指向 CompUnit
        }
        ;

FuncDef: type ID LPAREN param_list RPAREN Block
        {
            $$ = new ASTNode("FuncDef");
            $$->addChild($1);
            $$->addChild(new ASTNode("ID", *$2));
            $$->addChild($4);
            $$->addChild($6);
            delete $2;  // 释放 lexer 分配的字符串
        }
        ;

type: INT   { $$ = new ASTNode("int"); }
    | VOID  { $$ = new ASTNode("void"); }
    ;

param_list: /* empty */ { $$ = new ASTNode("ParamList"); }
          | Param { $$ = new ASTNode("ParamList"); $$->addChild($1); }
          | param_list COMMA Param { $1->addChild($3); $$ = $1; }
          ;

Param: INT ID 
     {
         $$ = new ASTNode("Param");
         $$->addChild(new ASTNode("ID", *$2));
         delete $2;  // 释放 lexer 分配的字符串
     }
     ;

Stmt: Block { $$ = $1; }
    | SEMICOLON { $$ = new ASTNode("EmptyStmt"); }
    | Expr SEMICOLON { $$ = $1; }
    | ID ASSIGN Expr SEMICOLON 
        { 
            $$ = new ASTNode("AssignStmt");
            $$->addChild(new ASTNode("ID", *$1));
            $$->addChild($3);
            delete $1;  // 释放 lexer 分配的字符串
        }
    | INT ID ASSIGN Expr SEMICOLON
        {
            $$ = new ASTNode("DeclAssignStmt");
            $$->addChild(new ASTNode("int"));
            $$->addChild(new ASTNode("ID", *$2));
            $$->addChild($4);
            delete $2;  // 释放 lexer 分配的字符串
        }
    | IF LPAREN Expr RPAREN Stmt
        {
            $$ = new ASTNode("IfStmt");
            $$->addChild($3);
            $$->addChild($5);
        }
    | IF LPAREN Expr RPAREN Stmt ELSE Stmt
        {
            $$ = new ASTNode("IfElseStmt");
            $$->addChild($3);
            $$->addChild($5);
            $$->addChild($7);
        }
    | WHILE LPAREN Expr RPAREN Stmt
        {
            $$ = new ASTNode("WhileStmt");
            $$->addChild($3);
            $$->addChild($5);
        }
    | BREAK SEMICOLON { $$ = new ASTNode("BreakStmt"); }
    | CONTINUE SEMICOLON { $$ = new ASTNode("ContinueStmt"); }
    | RETURN SEMICOLON {
        $$ = new ASTNode("ReturnStmt");
        $$->value = "void";
      }
    | RETURN Expr SEMICOLON {
        $$ = new ASTNode("ReturnStmt");
        $$->addChild($2);
        $$->value = "with_value";
      }
    ;

Block: LBRACE Stmt_list RBRACE
     {
         $$ = new ASTNode("Block");
         for (auto& child : $2->children) {
             $$->addChild(child.release());
         }
         delete $2;
     }
     ;

Stmt_list: /* empty */ { $$ = new ASTNode("StmtList"); }
         | Stmt_list Stmt { $1->addChild($2); $$ = $1; }
         ;

Expr: LOrExpr { $$ = $1; }
    ;
LOrExpr: LAndExpr { $$ = $1; }
       | LOrExpr OR LAndExpr 
       {
           $$ = new ASTNode("LOrExpr");
           $$->addChild($1);
           $$->addChild($3);
       }
       ;
LAndExpr: RelExpr { $$ = $1; }
        | LAndExpr AND RelExpr
        {
            $$ = new ASTNode("LAndExpr");
            $$->addChild($1);
            $$->addChild($3);
        }
        ;
RelExpr: AddExpr { $$ = $1; }
       | RelExpr LT AddExpr 
       {
           $$ = new ASTNode("RelExpr", "<");
           $$->addChild($1);
           $$->addChild($3);
       }
       | RelExpr GT AddExpr 
       {
           $$ = new ASTNode("RelExpr", ">");
           $$->addChild($1);
           $$->addChild($3);
       }
       | RelExpr LE AddExpr 
       {
           $$ = new ASTNode("RelExpr", "<=");
           $$->addChild($1);
           $$->addChild($3);
       }
       | RelExpr GE AddExpr 
       {
           $$ = new ASTNode("RelExpr", ">=");
           $$->addChild($1);
           $$->addChild($3);
       }
       | RelExpr EQ AddExpr 
       {
           $$ = new ASTNode("RelExpr", "==");
           $$->addChild($1);
           $$->addChild($3);
       }
       | RelExpr NE AddExpr 
       {
           $$ = new ASTNode("RelExpr", "!=");
           $$->addChild($1);
           $$->addChild($3);
       }
       ;

AddExpr: MulExpr { $$ = $1; }
       | AddExpr ADD MulExpr 
       {
           $$ = new ASTNode("AddExpr", "+");
           $$->addChild($1);
           $$->addChild($3);
       }
       | AddExpr SUB MulExpr 
       {
           $$ = new ASTNode("AddExpr", "-");
           $$->addChild($1);
           $$->addChild($3);
       }
       ;

MulExpr: UnaryExpr { $$ = $1; }
       | MulExpr MUL UnaryExpr 
       {
           $$ = new ASTNode("MulExpr", "*");
           $$->addChild($1);
           $$->addChild($3);
       }
       | MulExpr DIV UnaryExpr 
       {
           $$ = new ASTNode("MulExpr", "/");
           $$->addChild($1);
           $$->addChild($3);
       }
       | MulExpr MOD UnaryExpr 
       {
           $$ = new ASTNode("MulExpr", "%");
           $$->addChild($1);
           $$->addChild($3);
       }
       ;
UnaryExpr: PrimaryExpr { $$ = $1; }
         | ADD UnaryExpr 
         {
             $$ = new ASTNode("UnaryExpr", "+");
             $$->addChild($2);
         }
         | SUB UnaryExpr 
         {
             $$ = new ASTNode("UnaryExpr", "-");
             $$->addChild($2);
         }
         | NOT UnaryExpr 
         {
             $$ = new ASTNode("UnaryExpr", "!");
             $$->addChild($2);
         }
         ;

PrimaryExpr: ID 
           {
               $$ = new ASTNode("VarRef", *$1);
               delete $1;  // 释放 lexer 分配的字符串
           }
           | NUMBER 
           {
               $$ = new ASTNode("NumLiteral", std::to_string($1));
           }
           | LPAREN Expr RPAREN { $$ = $2; }
           | ID LPAREN arg_list RPAREN
           {
               $$ = new ASTNode("FuncCall", *$1);
               for (auto& child : $3->children) {
                   $$->addChild(child.release());
               }
               delete $1;  // 释放 lexer 分配的字符串
               delete $3;
           }
           ;

arg_list: /* empty */ { $$ = new ASTNode("ArgList"); }
        | Expr { $$ = new ASTNode("ArgList"); $$->addChild($1); }
        | arg_list COMMA Expr { $1->addChild($3); $$ = $1; }
        ;

%%
